#!/usr/bin/env python3
import cv2
import numpy as np
import pandas as pd
import xml.etree.ElementTree as ET

IMU_COLUMN_NAMES = ['lat', 'lon', 'alt', 'roll', 'pitch', 'yaw', 'vn', 've', 'vf', 'vl', 'vu', 'ax',
                    'ay', 'az', 'af', 'al', 'au', 'wx', 'wy', 'wz', 'wf', 'wl', 'wu', 'posacc', 'velacc',
                    'navstat', 'numsats', 'posmode', 'velmode', 'orimode']

def read_camera(path):
    return cv2.imread(path)

def read_point_cloud(path):
    return np.fromfile(path, dtype=np.float32).reshape(-1, 4)

def read_imu(path):
    df = pd.read_csv(path, header=None, sep=' ')
    df.columns = IMU_COLUMN_NAMES
    return df

def get_3d_boxes(tracklets, frame):
    """
    从tracklets中提取给定frame的boxes_3d的函数

    参数:
    tracklets (list): 包含tracklet数据的列表
    frame (int): 要提取boxes_3d的帧编号

    返回:
    list: 包含所有提取的boxes_3d的列表
    """
    boxes_3d = []
    for tracklet in tracklets:
        if frame >= tracklet.first_frame and frame < tracklet.first_frame + tracklet.num_frames:
            frame_idx = frame - tracklet.first_frame
            box_3d = {
                'h': tracklet.size[0],
                'w': tracklet.size[1],
                'l': tracklet.size[2],
                'tx': tracklet.trans[frame_idx, 0],
                'ty': tracklet.trans[frame_idx, 1],
                'tz': tracklet.trans[frame_idx, 2],
                'rz': tracklet.rot[frame_idx, 2]
            }
            boxes_3d.append(box_3d)
    return boxes_3d


def get_frame_object_types(tracklets):
    # 用于存储每一帧中的检测目标类别
    frame_data = []

    for tracklet in tracklets:
        object_type = tracklet.object_type
        start_frame = tracklet.first_frame
        end_frame = start_frame + tracklet.num_frames
        
        for frame in range(start_frame, end_frame):
            frame_data.append({'frame': frame, 'type': object_type})

    return pd.DataFrame(frame_data)

def get_frame_fields(tracklets):
    # 用于存储每一帧中的所有字段信息
    frame_data = []

    for tracklet_idx, tracklet in enumerate(tracklets):
        object_type = tracklet.object_type
        start_frame = tracklet.first_frame
        num_frames = tracklet.num_frames
        end_frame = start_frame + num_frames
        is_finished = True  # 默认情况下假设tracklet已完成解析
        has_amt = tracklet.amt_occ is not None and tracklet.amt_border is not None
        
        for frame in range(start_frame, end_frame):
            frame_info = {
                'frame': frame,
                'tracklet_idx': tracklet_idx,  # 添加tracklet索引
                'type': object_type,
                'h': tracklet.size[0],
                'w': tracklet.size[1],
                'l': tracklet.size[2],
                'tx': tracklet.trans[frame - start_frame, 0],
                'ty': tracklet.trans[frame - start_frame, 1],
                'tz': tracklet.trans[frame - start_frame, 2],
                'rx': tracklet.rot[frame - start_frame, 0],
                'ry': tracklet.rot[frame - start_frame, 1],
                'rz': tracklet.rot[frame - start_frame, 2],
                'state': tracklet.state[frame - start_frame],
                'occ_0': tracklet.occ[frame - start_frame, 0],
                'occ_kf': tracklet.occ[frame - start_frame, 1],
                'trunc': tracklet.trunc[frame - start_frame],
                'amt_occ_0': tracklet.amt_occ[frame - start_frame, 0] if tracklet.amt_occ is not None else None,
                'amt_occ_kf': tracklet.amt_occ[frame - start_frame, 1] if tracklet.amt_occ is not None else None,
                'amt_border_l': tracklet.amt_border[frame - start_frame, 0] if tracklet.amt_border is not None else None,
                'amt_border_r': tracklet.amt_border[frame - start_frame, 1] if tracklet.amt_border is not None else None,
                'amt_border_kf': tracklet.amt_border[frame - start_frame, 2] if tracklet.amt_border is not None else None,
                'is_finished': is_finished,
                'has_amt': has_amt
            }
            frame_data.append(frame_info)

    return pd.DataFrame(frame_data)



